import multiprocessing
try:
    import ogr
    import gdal
    import gdalconst
except:
    from osgeo import ogr
    from osgeo import gdal
    from osgeo import gdalconst
import numpy as np
import math
import os
import re


# clip raster file to tif file no same
def raster_clip(args):
	raster_path, save_folder, xcount, ycount = args
	raster_name = os.path.basename(raster_path)
	name = os.path.splitext(raster_name)[0]
	gdal.AllRegister()
	dataset = gdal.Open(raster_path)
	width,height = dataset.RasterXSize,dataset.RasterYSize
	projection = dataset.GetProjection()
	geotransform = dataset.GetGeoTransform()
	xmin, xsize = geotransform[0], geotransform[1]
	ymax, ysize = geotransform[3], geotransform[5]
	xstrid = width//xcount
	ystrid = height//ycount
	print(f"{raster_path}")
	for x in range(xcount):
		for y in range(ycount):	
			wxsize = xstrid if x!=xcount-1 else xstrid+width%xcount
			wysize = ystrid if y!=ycount-1 else ystrid+height%ycount
			left = xmin + xsize  * x * xstrid
			right = left + xsize * wxsize 
			top = ymax + ysize * y * ystrid
			bottom = top + ysize * wysize
			outputBounds = (left, bottom, right, top)
			path = os.path.join(save_folder, f"{name}_{x}_{y}.tif")
			if not os.path.exists(path):
				gdal.Warp(
					path, 
					raster_path, 
					format='GTiff', 
					outputBounds=outputBounds, 
					outputBoundsSRS=projection, 
					multithread=10)
	return


# clip raster by windows size
def window_clip_raster(args):
    raster_path, save_folder, window, step = args

    raster_name = os.path.basename(raster_path)
    name = os.path.splitext(raster_name)[0]

    gdal.AllRegister()
    dataset = gdal.Open(raster_path)
    width,height = dataset.RasterXSize,dataset.RasterYSize
    projection = dataset.GetProjection()
    geotransform = dataset.GetGeoTransform()
    xmin, xsize = geotransform[0], geotransform[1]
    ymax, ysize = geotransform[3], geotransform[5]

    xcount = math.ceil((width-window[0])/step[0])
    ycount = math.ceil((height-window[1])/step[1])

    paths = []
    for x in range(xcount):
        for y in range(ycount):	
            left = xmin + x * step[0] * xsize
            right = left + window[0] * xsize
            top = ymax + y * step[1] * ysize
            bottom = top + window[1] * ysize
            outputBounds = (left, bottom, right, top)
            path = os.path.join(save_folder, f"{name}_{x}_{y}.tif")
            if not os.path.exists(path):
                gdal.Warp(
                    path, 
                    raster_path, 
                    format='GTiff', 
                    outputBounds=outputBounds, 
                    outputBoundsSRS=projection, 
                    multithread=10)
            paths.append(path)
    return paths


# vector transform raster by model raster
def vector_2_raster(args):
    [vector_path, raster_path, result_path] = args
    if os.path.exists(result_path):
        return

    bands=[1]
    burn_values=[1]
    field=""
    all_touch="False"
    # NoData_value = 255

    gdal.UseExceptions()
    ogr.RegisterAll()
    gdal.SetConfigOption("SHAPE_ENCODING", "GBK")
    datasource = ogr.Open(vector_path)
    layer = datasource.GetLayer()

    gdal.AllRegister()
    dataset = gdal.Open(raster_path, gdalconst.GA_ReadOnly)
    width,height = dataset.RasterXSize,dataset.RasterYSize
    projection = dataset.GetProjection()
    geotransform = dataset.GetGeoTransform()

    driver = gdal.GetDriverByName('GTiff')
    target_dataset = driver.Create(result_path,width,height,1,gdal.GDT_Byte)
    target_dataset.SetGeoTransform(geotransform)
    target_dataset.SetProjection(projection)
    band = target_dataset.GetRasterBand(1)
    # band.SetNoDataValue(NoData_value)
    band.FlushCache()
    gdal.RasterizeLayer(
        target_dataset, 
        bands,
        layer, 
        burn_values=burn_values,
        options=["ALL_TOUCHED="+all_touch,"ATTRIBUTE="+field] if field else ["ALL_TOUCHED="+all_touch]
    )
    band = target_dataset.GetRasterBand(1)
    array = band.ReadAsArray()
    target_dataset.FlushCache()
    dataset.FlushCache()
    del dataset
    del target_dataset
    if not(np.any(array)):
        os.remove(result_path)
        os.remove(raster_path)
    return


# start
def main_1():
    raster_folder = r"F:\GoogleImage\36_JiangXi\2022\赣州市\已做完的\提取raster文件"
    vector_path = r"E:\赣州市_柑橘第一批shp\orange_赣州市.shp"
    image_folder = r"G:\DATA\images"
    lable_folder = r"G:\DATA\labels"
    window, step = [512,512], [512,512]

    print("big raster clip small tif")
    args_list = []
    for file_name in os.listdir(raster_folder):
        if re.match(".*\.tif$|.*\.jp2$",file_name):
            raster_path = os.path.join(raster_folder, file_name)
            args_list.append([raster_path, image_folder, window, step])  
    pool = multiprocessing.Pool(2)
    resluts = pool.map(window_clip_raster, args_list)
    pool.close()
    pool.join()
    
    print("get label raster")
    args_list = []
    for tif_paths in resluts:
        for tif_path in tif_paths:
            tif_name = os.path.basename(tif_path)
            result_path = os.path.join(lable_folder, tif_name)
            args_list.append([vector_path, tif_path, result_path])
            # vector_2_raster([vector_path, tif_path, result_path])
    print(len(args_list))

    pool = multiprocessing.Pool(2)
    resluts = pool.map(vector_2_raster, args_list)
    pool.close()
    pool.join()
    return


def task(args):

    raster_path, vector_path, image_folder, window, step, lable_folder = args

    raster_name = os.path.basename(raster_path)
    name = os.path.splitext(raster_name)[0]

    gdal.AllRegister()
    dataset = gdal.Open(raster_path)
    width,height = dataset.RasterXSize,dataset.RasterYSize
    projection = dataset.GetProjection()
    geotransform = dataset.GetGeoTransform()
    xmin, xsize = geotransform[0], geotransform[1]
    ymax, ysize = geotransform[3], geotransform[5]

    xcount = math.ceil((width-window[0])/step[0])
    ycount = math.ceil((height-window[1])/step[1])

    for x in range(xcount):
        for y in range(ycount):	
            left = xmin + x * step[0] * xsize
            right = left + window[0] * xsize
            top = ymax + y * step[1] * ysize
            bottom = top + window[1] * ysize
            outputBounds = (left, bottom, right, top)
            path = os.path.join(image_folder, f"{name}_{x}_{y}.tif")
            if not os.path.exists(path):
                gdal.Warp(
                    path, 
                    raster_path, 
                    format='GTiff', 
                    outputBounds=outputBounds, 
                    outputBoundsSRS=projection, 
                    multithread=10)
            tif_name = os.path.basename(path)
            result_path = os.path.join(lable_folder, tif_name)

            vector_2_raster([vector_path, path, result_path])

    return


def main():
    raster_folder = r"F:\GoogleImage\36_JiangXi\2022\赣州市\已做完的\提取raster文件"
    vector_path = r"E:\赣州市_柑橘第一批shp\orange_赣州市.shp"
    image_folder = r"G:\DATA1\images"
    lable_folder = r"G:\DATA1\labels"
    window, step = [512,512], [512,512]

    args_list = []
    for file_name in os.listdir(raster_folder):
        if re.match(".*\.tif$|.*\.jp2$",file_name):
            raster_path = os.path.join(raster_folder, file_name)
            args_list.append([raster_path, vector_path, image_folder, window, step, lable_folder])  
            # task([raster_path, vector_path, image_folder, window, step, lable_folder])

    pool = multiprocessing.Pool(4)
    pool.map(task, args_list)
    pool.close()
    pool.join()
    return

if __name__ == "__main__":
    main()


    